package com.tfswx.ruleengine.compute.util.sm4;

import java.io.FilterOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.concurrent.atomic.AtomicInteger;

public abstract class DecorateOutputStreamBase extends FilterOutputStream {

    private final static int STATUS_INIT = 0;
    private final static int STATUS_WRITING = 1;
    private final static int STATUS_CLOSED = 2;

    private int writeSize;
    private int destSize;
    private byte[] writeBuffer;
    private byte[] destBuffer;
    private final byte[] singletonBuffer = new byte[1];

    private DecorateOutputStreamBase(OutputStream out) {
        super(out);
    }

    protected DecorateOutputStreamBase(OutputStream out, int writeSize, int destSize) {
        this(out);

        this.writeSize = writeSize;
        this.destSize = destSize;

        innerInit();
    }

    private volatile int writeLength;

    /**
     * 当前流状态
     */
    private final AtomicInteger status = new AtomicInteger(0);

    /**
     * 初始化最后执行
     */
    protected void innerInit() {

    }

    /**
     * 流开始读取处理方法
     */
    protected void innerOpen() {

    }

    /**
     * 流结束读取处理方法
     */
    protected void innerClose() {

    }

    /**
     * 内部数据真实处理方法
     * (该方法执行完成时，sourceBuffer 数据将全部清空)
     *
     * @param sourceBuffer
     * @param sourceLength
     * @param destBuffer
     * @return destBuffer 缓冲区内真实数据长度，返回 -1 时表示没有数据
     */
    protected abstract int innerWriteFrom(byte[] sourceBuffer, int sourceLength, byte[] destBuffer);

    /**
     * 开始写入
     */
    private void initStart() throws IOException {
        if (status.compareAndSet(STATUS_INIT, STATUS_WRITING)) {

            if (writeBuffer == null) {
                writeBuffer = new byte[writeSize];
            }
            if (destBuffer == null) {
                destBuffer = new byte[destSize];
            }

            innerOpen();
        } else if (status.get() == STATUS_CLOSED) {
            throw new IOException("stream is closed.");
        }
    }

    @Override
    public void write(int b) throws IOException {

        initStart();

        singletonBuffer[0] = (byte) b;
        this.write(singletonBuffer, 0, 1);
    }

    @Override
    public void write(byte[] b) throws IOException {
        this.write(b, 0, b.length);
    }

    @Override
    public void write(byte[] b, int off, int len) throws IOException {

        initStart();

        int inputOffset = off;
        int inputLength = len;

        int leastCount, destCount;
        do {
            leastCount = writeSize - writeLength;
            if (leastCount > inputLength) {
                // 数据不够，当前写入缓冲区后就立刻中止，等待下次新数据来了继续操作
                System.arraycopy(b, inputOffset, writeBuffer, writeLength, inputLength);
                writeLength += inputLength;
                return;
            } else {
                // 数据够，读取完立即操作处理，然后继续下一段数据操作
                System.arraycopy(b, inputOffset, writeBuffer, writeLength, leastCount);
                destCount = innerWriteFrom(writeBuffer, writeSize, destBuffer);
                writeLength = 0; // 清空已处理数据
                inputLength -= leastCount; // 源数据消耗
                inputOffset += leastCount; // 读取位置后移
                if (destCount <= 0) {
                    continue;
                }
                out.write(destBuffer, 0, destCount);
            }
        } while (inputLength > 0);
    }

    @Override
    public void close() throws IOException {

        if (status.compareAndSet(STATUS_WRITING, STATUS_CLOSED)) {

            // 操作最后一段数据
            if (writeLength > 0) {
                int destCount = innerWriteFrom(writeBuffer, writeLength, destBuffer);
                writeLength = -1;
                if (destCount > 0) {
                    out.write(destBuffer, 0, destCount);
                }
            } else if (writeLength == 0) {
                int destCount = innerWriteFrom(writeBuffer, 0, destBuffer);
                writeLength = -1;
                if (destCount > 0) {
                    out.write(destBuffer, 0, destCount);
                }
            }

            innerClose();
        } else {
            status.set(STATUS_CLOSED);
        }

        super.close();
    }
}
